/***************************************************************************************************
 * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights
 *reserved. SPDX-License-Identifier: BSD-3-Clause
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 * this list of conditions and the following disclaimer in the documentation
 * and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the copyright holder nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *POSSIBILITY OF SUCH DAMAGE.
 *
 **************************************************************************************************/
#pragma once

#include "asymmetric/epilogue/threadblock/default_epilogue_tensor_op_dequant.h"
#include "asymmetric/gemm/kernel/sparse_gemm_dequant.h"
#include "cutlass/gemm/kernel/default_gemm_sparse.h"

////////////////////////////////////////////////////////////////////////////////

namespace cutlass {
namespace gemm {
namespace kernel {
namespace asymmetric {

////////////////////////////////////////////////////////////////////////////////

template <
    /// Element type for A matrix operand
    typename ElementA_,
    /// Layout type for A matrix operand
    typename LayoutA_,
    /// Access granularity of A matrix in units of elements
    int kAlignmentA,
    /// Element type for B matrix operand
    typename ElementB_,
    /// Layout type for B matrix operand
    typename LayoutB_,
    /// Access granularity of B matrix in units of elements
    int kAlignmentB,
    /// Element type for C and D matrix operands
    typename ElementC_,
    /// Layout type for C and D matrix operands
    typename LayoutC_,
    /// Element type for internal accumulation
    typename ElementAccumulator,
    /// Operator class tag
    typename OperatorClass,
    /// Tag indicating architecture to tune for
    typename ArchTag,
    /// Threadblock-level tile size (concept: GemmShape)
    typename ThreadblockShape,
    /// Warp-level tile size (concept: GemmShape)
    typename WarpShape,
    /// Warp-level tile size (concept: GemmShape)
    typename InstructionShape,
    /// Epilogue output operator
    typename EpilogueOutputOp,
    /// Threadblock-level swizzling operator
    typename ThreadblockSwizzle,
    /// Number of stages used in the pipelined mainloop
    int Stages,
    /// If true, kernel is configured to support serial reduction in the
    /// epilogue
    bool SplitKSerial,
    /// Operation performed by GEMM
    typename Operator>
struct DefaultSparseGemmDequant
    : public DefaultSparseGemm<
          ElementA_, LayoutA_, kAlignmentA, ElementB_, LayoutB_, kAlignmentB,
          ElementC_, layout::RowMajor, ElementAccumulator,
          arch::OpClassTensorOp, arch::Sm80, ThreadblockShape, WarpShape,
          InstructionShape, EpilogueOutputOp, ThreadblockSwizzle, Stages,
          SplitKSerial, Operator> {
  using DefaultSparseGemm = DefaultSparseGemm<
      ElementA_, LayoutA_, kAlignmentA, ElementB_, LayoutB_, kAlignmentB,
      ElementC_, layout::RowMajor, ElementAccumulator, arch::OpClassTensorOp,
      arch::Sm80, ThreadblockShape, WarpShape, InstructionShape,
      EpilogueOutputOp, ThreadblockSwizzle, Stages, SplitKSerial, Operator>;

  using Epilogue = typename cutlass::epilogue::threadblock::asymmetric::
      DefaultEpilogueTensorOpDequant<
          ThreadblockShape, typename DefaultSparseGemm::Mma::Operator,
          DefaultSparseGemm::kPartitionsK, EpilogueOutputOp,
          EpilogueOutputOp::kCount>::Epilogue;

  /// Define the kernel-level GEMM operator.
  using GemmKernel =
      kernel::asymmetric::SparseGemmDequant<typename DefaultSparseGemm::Mma,
                                            Epilogue, ThreadblockSwizzle,
                                            SplitKSerial>;
};

////////////////////////////////////////////////////////////////////////////////

}  // namespace asymmetric
}  // namespace kernel
}  // namespace gemm
}  // namespace cutlass
